import * as tfvis from '@tensorflow/tfjs-vis'
import * as tf from '@tensorflow/tfjs'
import { getData } from './data.js'

/**
 * XOR
 */
window.onload = () => {
	const data = getData(400)

	// 散点图
	tfvis.render.scatterplot(
	{
		name: "XOR 训练集"
	}, {
		values: [
			data.filter(p => p.label === 1),
			data.filter(p => p.label === 0)
		]
	}, {
		xAxisDomain: [-6, 6],
		yAxisDomain: [-6, 6]
	})

	/**
	 * 连续的模型
	 */
	const model = tf.sequential()

	/** 隐藏全链接层
	 * units 神经元数
	 * inputShape 输入形状
	 */
	model.add(tf.layers.dense({
		units: 4,
		inputShape: [2],
		activation: 'relu'
	}))

	/** 输出全链接层
	 * units 神经元数
	 */
	model.add(tf.layers.dense({
		units: 1,
		activation: 'sigmoid'
	}))

	/**
	 * 设置损失函数
	 * 优化器adam(学习率)
	 */
	model.compile({
		loss: tf.losses.logLoss,
		optimizer: tf.train.adam(0.1)
	})

	const inputs = tf.tensor(data.map(p => [p.x, p.y]));
	const labels = tf.tensor(data.map(p => p.label));

	/** fit 开始学习任务
	 * batchSize 取样数
	 * epochs 对训练数据数组进行迭代的次数。
	 */
	model.fit(inputs, labels, {
		batchSize: 40,
		epochs: 20,
		callbacks: tfvis.show.fitCallbacks({
			name: "训练过程"
		}, [
			'loss'
		])
	})

	window.predict = (form) => {
		const output = model.predict(tf.tensor([[form.x.value * 1, form.y.value * 1]]));
		alert(`如果 x,y 为 ${form.x.value} | ${form.y.value}，那么预测为 ${output.dataSync()[0]}`);
	}
}

